# coding: utf-8

from __future__ import absolute_import, division, print_function, unicode_literals

import os
from pathlib import Path

from pypeg2 import compose
from thrift2pyi.peg import Consts, Enums, Exceptions, Imports, Services, Unions
from yapf.yapflib.yapf_api import FormatCode

from apitest.cmd.peg import Init, PYI, Struct, Structs
from apitest.cmd.thrift2pyi import Thrift2pyi


class Thrift2py(Thrift2pyi):

    def __init__(self, filename, prefix, out, need_service=False):
        super().__init__(filename, prefix, out)
        self.pyi = PYI()
        self.pyi.imports = Imports()
        self.pyi.services = Services()
        self.pyi.structs = Structs()
        self.pyi.unions = Unions()
        self.pyi.enums = Enums()
        self.pyi.exceptions = Exceptions()
        self.pyi.consts = Consts()
        self.need_service = need_service

    def _spec2type(self, spec):
        _type = super()._spec2type(spec)
        if isinstance(spec[-1], bool) and spec[-1] is False:
            self._imports["typing"].add("Optional")
            return "Optional[%s]" % _type
        return _type

    def _struct2pyi(self, struct):
        self._imports["pydantic"].add("BaseModel")
        _struct = super()._struct2pyi(struct)
        p_struct = Struct()
        p_struct.name = _struct.name
        p_struct.annotations = _struct.annotations
        p_init = Init()
        p_init.params = _struct.init.params
        p_struct.init = p_init
        return p_struct

    def _structs2pyi(self):
        structs = []
        for struct in self.meta["structs"]:
            structs.append(self._struct2pyi(struct))
        for struct in structs:
            self.pyi.structs.append(struct)

    def _scan_includes(self):
        for include in self.meta["includes"]:
            thrift_file = include.__thrift_file__
            module = thrift_file.split(os.sep)[-1].split(".")[0]
            rel_path = os.path.relpath(include.__thrift_file__, os.path.dirname(self.thrift.__thrift_file__))
            rel_dir = os.path.dirname(rel_path)
            self._module2package[module] = "." + rel_dir.replace(".." + os.sep, "..").replace("." + os.sep, ".")

            if not os.path.exists(rel_path):
                out = os.path.join(Path(self.out).resolve().parent,
                                   "{}.{}".format(rel_path.replace(".", "_"), self.out.rsplit(".", 1)[-1]))
                Thrift2py(include.__thrift_file__, self.prefix, out).output()

    def _thrift2pyi(self):
        self._imports["__future__"].add("annotations")
        self._consts2pyi()
        self._scan_includes()
        self._structs2pyi()
        self._unions2pyi()
        self._excs2pyi()
        if self.need_service:
            self._services2pyi()
        self._enums2pyi()
        self._includes2pyi()

    def output(self):
        self._thrift2pyi()
        try:
            o = FormatCode(compose(self.pyi))[0]
        except:
            o = compose(self.pyi)
            print("format failed")
        if not self.out or self.prefix == self.out:
            with open("%s.py" % self.filename.replace(".thrift", ""), "w") as f:
                f.write(o)
            return

        with open("%s" % self.out, "w") as f:
            f.write(o)
